"""
小米发布的预训练GemmaX2模型，支持东盟语言的翻译。
支持的语言：阿拉伯语、孟加拉语、捷克语、德语、英语、西班牙语、波斯语、法语、希伯来语、印地语、印度尼西亚语、意大利语、日语、
柬埔寨语、韩语、老挝语、马来语、缅甸语、荷兰语、波兰语、葡萄牙语、俄语、泰语、他加禄语、土耳其语、乌尔都语、越南语、中文。

运行环境:
    Ubuntu 20.04.5 LTS + python:3.8
安装包：
    pip install transformers
下载模型：
    https://hf-mirror.com/ModelSpace/GemmaX2-28-9B-v0.1
    from modelscope import snapshot_download
    snapshot_download('ModelSpace/GemmaX2-28-9B-v0.1', local_dir=r'/root/GemmaX2-28-9B-v0.1')
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
from translate.languages import xiaomi_languages
from threading import RLock
import re
import torch


class XiaomiTranslator(object):

    lock = RLock()

    def __new__(cls, *args, **kwargs):
        with XiaomiTranslator.lock:
            if not hasattr(XiaomiTranslator, "_instance"):
                XiaomiTranslator._instance = object.__new__(cls)
                cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
                repo_id = "/root/autodl-tmp/GemmaX2-28-9B-v0.1"
                cls.tokenizer = AutoTokenizer.from_pretrained(repo_id)
                cls.model = AutoModelForCausalLM.from_pretrained(repo_id, 
                torch_dtype=torch.float16, 
                trust_remote_code=True, device_map=cls.device)
        return XiaomiTranslator._instance

    def translate(self, text: str, from_lan: str='中文', to_lan='泰语'):
        """翻译"""
        from_lan = xiaomi_languages.get(from_lan, from_lan)
        to_lan = xiaomi_languages.get(to_lan, to_lan)
        
        if from_lan == to_lan:
            return text
        prompt = f"""Translate this from {from_lan} to {to_lan}:\n{from_lan}:{text}\n{to_lan}:"""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        outputs = self.model.generate(**inputs, max_new_tokens=512)
        out_text = self.tokenizer.decode(outputs[0], skip_special_tokens=False)
        return self.extract_result(out_text, from_lan, to_lan)

    @staticmethod
    def extract_result(out_text: str, from_lan: str='Chinese', to_lan='Thai'):
        """从LLM的输出结果抽取翻译文本"""
        from_lan = xiaomi_languages.get(from_lan, from_lan)
        to_lan = xiaomi_languages.get(to_lan, to_lan)
        founds = re.findall(f'{to_lan}:\s*([^\n]*)\s*?', out_text, re.DOTALL)
        trans_texts = list(filter(lambda item: item.find(from_lan) == -1, founds))
        if len(trans_texts) == 0:
            return out_text
        longest_text = ''
        for item in trans_texts:
            if len(item) > len(longest_text):
                longest_text = item
        return re.sub('<[^<>]+>$', '', longest_text)


if __name__ == '__main__':
    xiaomi = XiaomiTranslator()
    trans_text = xiaomi.translate(text="敏感词汇怎么屏蔽?", from_lan='中文', to_lan='老挝语')
    print(trans_text)